import sqlite3
import csv

# 连接数据库
conn = sqlite3.connect("E:\\immunedb\\immuno.db")
cursor = conn.cursor()

# 目标HLA(除了这个示例其他hla都要跑）
target_hlas = ["HLA-A*02:01", "HLA-A*01:01", "HLA-B*48:01"]

# 执行查询（直接拼接SQL，避免占位符可能的问题）
query = f"""
SELECT 
    p.hla,
    p.peptide_id,
    g.geneName,
    g.avg_exp,
    e.exp_level
FROM 
    predictions p
JOIN peptides pt ON p.peptide_id = pt.peptide_id
JOIN protein_gene g ON pt.protein_id = g.protein_id
LEFT JOIN protein_exp e ON g.geneName = e.geneName
WHERE 
    p.hla IN ('{"', '".join(target_hlas)}')
    AND p.probability > 0.5
"""

try:
    cursor.execute(query)
    rows = cursor.fetchall()
    print(f"成功获取 {len(rows)} 条记录")
except Exception as e:
    print(f"查询执行错误：{str(e)}")
    rows = []

conn.close()

# 初始化统计字典
stats = {
    hla: {
        "only_protein": 0,
        "only_rna": 0,
        "both": 0,
        "neither": 0,
        "total": 0
    } for hla in target_hlas
}

# 处理每条记录
for row in rows:
    try:
        hla = row[0]
        avg_exp = row[3]
        exp_level = row[4]
        
        # 判断RNA表达
        rna_exp = False
        if avg_exp is not None:
            try:
                avg_exp_num = float(avg_exp)
                if avg_exp_num != 0:
                    rna_exp = True
            except:
                pass
        
        # 判断蛋白表达
        protein_exp = False
        if exp_level is not None:
            if exp_level.strip() in ["High", "Low", "Medium"]:
                protein_exp = True
        
        # 更新统计
        stats[hla]["total"] += 1
        if rna_exp and protein_exp:
            stats[hla]["both"] += 1
        elif rna_exp:
            stats[hla]["only_rna"] += 1
        elif protein_exp:
            stats[hla]["only_protein"] += 1
        else:
            stats[hla]["neither"] += 1
    except Exception as e:
        print(f"处理记录时出错：{str(e)}，跳过该记录")

# 写入CSV
with open("HLA表达统计结果.csv", "w", newline="", encoding="utf-8") as f:
    writer = csv.writer(f)
    writer.writerow(["HLA", "仅蛋白表达", "仅RNA表达", "RNA和蛋白都有表达", "RNA和蛋白都无表达", "总记录数"])
    for hla in target_hlas:
        writer.writerow([
            hla,
            stats[hla]["only_protein"],
            stats[hla]["only_rna"],
            stats[hla]["both"],
            stats[hla]["neither"],
            stats[hla]["total"]
        ])

print("统计完成，结果已写入 HLA表达统计结果.csv")